import json
import pickle

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report, \
	confusion_matrix
from sklearn.model_selection import train_test_split
from tqdm import tqdm

from __00__config import Config


def train_rf_model():
	# 获取训练数据
	train_data = pd.read_csv(config.process_train_datapath, sep='\t', encoding='utf-8')
	# 查看数据情况
	# print(train_data.info())
	# print(train_data.head())

	words = train_data['words']
	labels = train_data['labels']
	# 查看数据
	# print(words.head())
	# print(labels.head())

	# 数据预处理
	stop_words = [line.strip() for line in open(config.stop_words_path, encoding='utf-8').readlines()]
	# print(len(stop_word))

	# 实例化TfidfVectorizer
	tfidf = TfidfVectorizer(stop_words=stop_words)
	# 将文本转换为词频矩阵
	features = tfidf.fit_transform(words)
	# print(features.shape)  # (142145, 59763)
	# print(features)  # 查看tfidf训练后特征效果
	# # 转化为矩阵形式展示
	# # 打印特征矩阵，将稀疏矩阵转换为密集格式并输出
	# print(features.toarray())
	# # 打印TF-IDF模型中的所有特征名称列表
	# print(list(tfidf.get_feature_names_out()))
	# # 打印特征名称的数量
	# print(len(tfidf.get_feature_names_out()))  # 7910
	# # 打印TF-IDF模型的词汇表，包含所有特征及其对应的索引
	# print(tfidf.vocabulary_)
	# # 打印词汇表的大小，即特征数量
	# print(len(tfidf.vocabulary_))  # 7910

	# 划分训练集和测试集
	x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.2, random_state=42)
	# 实例化随机森林模型
	rf = RandomForestClassifier()
	# 训练模型
	print('开始训练模型...')
	for _ in tqdm(range(1), desc='训练进度'):
		rf.fit(x_train, y_train)
	print('开始模型预测和评估...')
	y_predict = rf.predict(x_test)
	# 模型评估
	print('模型评估结果如下：')
	print(f'准确率：{accuracy_score(y_test, y_predict)}')
	print(f'精确率：{precision_score(y_test, y_predict, average="macro")}')
	print(f'召回率：{recall_score(y_test, y_predict, average="macro")}')
	print(f'F1值：{f1_score(y_test, y_predict, average="macro")}')

	# 打印评估报告
	print(f'评估报告：\n{classification_report(y_test, y_predict)}')
	# 打印混淆矩阵
	print(f'混淆矩阵：\n{confusion_matrix(y_test, y_predict)}')

	print('开始保存模型和向量化器...')
	# 保存模型
	with open(config.rf_model_save_path, 'wb') as f:
		pickle.dump(rf, f)
	print(f'模型保存成功，路径为：{config.rf_model_save_path}')

	# 保存向量化器
	with open(config.tfidf_model_save_path, 'wb') as f:
		pickle.dump(tfidf, f)
	print(f'向量化器已保存，路径为：{config.tfidf_model_save_path}')


if __name__ == '__main__':
	config = Config()
	train_rf_model()
